package com.object.module.lq.sys.apiController;

import com.alibaba.dashscope.aigc.conversation.ConversationParam.ResultFormat;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.tools.*;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.utils.JsonUtils;
import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.github.victools.jsonschema.generator.*;
import com.object.model.*;
import com.object.model.execute.ModelExecute;
import com.object.module.lq.sys.entity.bean.TModel;
import com.object.module.lq.sys.entity.bean.TModelConfig;
import com.object.module.lq.sys.service.TModelConfigService;
import com.object.module.lq.sys.service.TModelService;
import io.reactivex.Flowable;
import lombok.SneakyThrows;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang.StringUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletResponse;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author Li Qing
 * date: 2025/6/4 17:15
 * description:        ai tool https://docs.spring.io/spring-ai/reference/api/tools.html
 * 同义 ai tool  https://help.aliyun.com/zh/model-studio/qwen-function-calling#809cfd466dvb4
 */
@RestController
public class ModelController {

    @Value("${model.apiKey}")
    private String apiKey;

    @Resource
    private TModelConfigService modelConfigService;


    @Resource
    private TModelService modelService;
    @Resource
    private Map<String, ModelExecute> modelExecuteMap;


    /**
     * ai 接口
     *
     * @param messages
     * @param
     */
    @PostMapping(value = "/modelApi", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    @SneakyThrows
    public void senModel(@RequestBody List<Message> messages, HttpServletResponse servletResponse) {
        // 设置响应类型为文本流 + UTF-8 编码
        servletResponse.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE);
        servletResponse.setCharacterEncoding(StandardCharsets.UTF_8.name());
        servletResponse.setHeader("Cache-Control", "no-cache");
        servletResponse.setHeader("X-Accel-Buffering", "no");
        // 关键：Nginx 不要缓冲

        Generation gen = new Generation();

        GenerationParam param = buildGenerationParam(drawSysMessage(messages));
        Flowable<GenerationResult> result = gen.streamCall(param);
        PrintWriter writer = servletResponse.getWriter();

        sendApi(messages, result, writer, gen);

    }

    private void sendApi(List<Message> messages, Flowable<GenerationResult> result, PrintWriter writer, Generation gen) throws NoApiKeyException, InputRequiredException {
        AtomicReference<String> toolCallId = new AtomicReference<>();
        // 用于记录当前处理的函数名、参数、调用 ID
        AtomicReference<String> currentFunctionName = new AtomicReference<>();

        AtomicReference<List<Message>> finalMessages = new AtomicReference<>();

        AtomicReference<String> accumulatedArguments = new AtomicReference<>("");


        result.blockingForEach(message -> {
            Message assistantMessage = message.getOutput().getChoices().get(0).getMessage();

            if (assistantMessage.getToolCalls() != null) {
                for (ToolCallBase toolCall : assistantMessage.getToolCalls()) {
                    if ("function".equals(toolCall.getType())) {
                        ToolCallFunction functionCall = (ToolCallFunction) toolCall;
                        String functionName = functionCall.getFunction().getName();
                        // 如果是新的函数调用，清空之前的缓存
                        if (!Objects.equals(functionName, currentFunctionName.get()) && functionName != null) {
                            currentFunctionName.set(functionName);
                        }
                        // 拼接 arguments
                        String currentArgument = functionCall.getFunction().getArguments();
                        if (currentArgument != null) {
                            accumulatedArguments.updateAndGet(s -> (s == null ? "" : s) + currentArgument);
                        }
                        // 如果 arguments 已经闭合（即包含完整的 JSON），才执行调用本地函数
                        sendTool(messages, accumulatedArguments, currentFunctionName, functionCall, finalMessages, toolCallId);
                    }
                }
            } else {
                String content = assistantMessage.getContent();
                if (content != null && !content.isEmpty()) {
                    writer.write(content);
                    writer.flush();
                }
            }
        });

        // 如果有 tool_call，再次调用模型生成最终回复
        if (finalMessages.get() != null) {
            GenerationParam nextParam = buildGenerationParam(finalMessages.get());
            Flowable<GenerationResult> finalResult = gen.streamCall(nextParam);

            finalResult.blockingForEach(message -> {
                String content = message.getOutput().getChoices().get(0).getMessage().getContent();
                writer.write(content);
                writer.flush();
            });
        }

        writer.close();
    }

    /**
     * 调用工具
     *
     * @param messages
     * @param accumulatedArguments
     * @param currentFunctionName
     * @param functionCall
     * @param finalMessages
     * @param toolCallId
     */
    private void sendTool(List<Message> messages, AtomicReference<String> accumulatedArguments, AtomicReference<String> currentFunctionName, ToolCallFunction functionCall, AtomicReference<List<Message>> finalMessages, AtomicReference<String> toolCallId) {
        if (isJsonComplete(accumulatedArguments.get())) {

            String fullArguments = accumulatedArguments.get();

            ModelExecute execute = modelExecuteMap.getOrDefault(currentFunctionName.get(), null);
            if (execute != null) {
                String resultContent = execute.execute(fullArguments);

                // 构造 tool_result 消息
                Message toolResultMessage = Message.builder()
                        .role("tool")
                        .content(resultContent)
                        .toolCallId(functionCall.getId())
                        .build();
                ToolCallFunction function1 = new ToolCallFunction();
                function1.setType("function");

                function1.setId(functionCall.getId());
                // 再创建内部类实例
                ToolCallFunction.CallFunction callFunction = function1.new CallFunction();

                callFunction.setName(currentFunctionName.get());
                callFunction.setArguments(accumulatedArguments.get());
                function1.setFunction(callFunction);
                Message toolCallMessage = Message.builder()
                        .role("assistant")
                        .toolCalls(Collections.singletonList(
                                function1
                        ))
                        .build();
                List<Message> newMessages = new ArrayList<>(messages);
                newMessages.add(toolCallMessage);
                newMessages.add(toolResultMessage);


                finalMessages.set(newMessages);
                toolCallId.set(functionCall.getId());

                // 清空缓存，避免重复执行
                accumulatedArguments.set("");
            }
        }
    }

    public static boolean isJsonComplete(String json) {
        if (json == null || json.trim().isEmpty()) {
            return false;
        }
        try {
            JSON.parse(json);
            // FastJSON 解析
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    private GenerationParam buildGenerationParam(List<com.alibaba.dashscope.common.Message> list) {

        /**
         * 加载ai 工具
         */
        List<ToolBase> listFunctions = new ArrayList<>();
        Arrays.stream(ModelConfigTool.values()).forEach(modelConfigTool -> {
            ObjectNode jsonNodes = generateSchema(modelConfigTool.getDtoClazz() != null ? modelConfigTool.getDtoClazz() : modelConfigTool.getClazz());
            FunctionDefinition fdWeather = FunctionDefinition.builder().name(modelConfigTool.getBeanName())
                    .description(modelConfigTool.getDescription())
                    .parameters(JsonUtils.parseString(jsonNodes.toString()).getAsJsonObject()).build();
            listFunctions.add(ToolFunction.builder().function(fdWeather).build());
        });


        return GenerationParam.builder()
                // 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(apiKey)
                .model(modelConfigService.getOne(new LambdaQueryWrapper<TModelConfig>()
                        .eq(TModelConfig::getModeStatus, true)
                ).getModelKey())
                // 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .messages(list)
                .incrementalOutput(true)
                .resultFormat(ResultFormat.MESSAGE)
                .tools(listFunctions)
                .build();
    }

    private static ObjectNode generateSchema(Class<?> clazz) {
        SchemaGeneratorConfigBuilder configBuilder =
                new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON);
        SchemaGeneratorConfig config = configBuilder.with(Option.EXTRA_OPEN_API_FORMAT_VALUES)
                .without(Option.FLATTENED_ENUMS_FROM_TOSTRING).build();
        SchemaGenerator generator = new SchemaGenerator(config);
        return generator.generateSchema(clazz);
    }

    private List<Message> drawSysMessage(List<Message> messages) {
        List<Message> msgList = new ArrayList<>();
        List<com.object.module.lq.sys.entity.bean.TModel> list = modelService.list();
        if (CollectionUtils.isNotEmpty(list)) {
            list.stream().forEach(i -> {
                if (StringUtils.isNotEmpty(i.getContent())) {
                    com.alibaba.dashscope.common.Message message = new com.alibaba.dashscope.common.Message();
                    message.setRole(ModelFunction.system.name());
                    message.setContent(i.getContent());
                    msgList.add(message);
                }
            });
        }
        messages.stream().forEach(i -> {
            if (StringUtils.isNotEmpty(i.getContent())) {
                com.alibaba.dashscope.common.Message message = new com.alibaba.dashscope.common.Message();
                message.setRole(i.getRole());
                message.setContent(i.getContent());
                msgList.add(message);
            }
        });
        return msgList;
    }
}
